﻿//Copyright (c) 2010 Christopher E. S. King (http://chriskingconsulting.com)
//
//Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files 
// (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, 
// publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do 
// so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 
// MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE 
// FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
// WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Reflection;
using System.Collections;
using System.Threading;

namespace King.Extensions {

    public static class TreeTraverse {

        public static IEnumerable<TResult> PreOrder<T, TResult>(
            T node,
            Func<T, IEnumerable<T>> getChildren,
            Func<T, int, TResult> visitor) {

            return PreOrPostOrder(node, getChildren, true, visitor);
        }
        public static IEnumerable<TResult> PostOrder<T, TResult>(
            T node,
            Func<T, IEnumerable<T>> getChildren,
            Func<T, int, TResult> visitor) {

            return PreOrPostOrder(node, getChildren, false, visitor);
        }
        public static IEnumerable<TResult> PreOrPostOrder<T, TResult>(
            T node,
            Func<T, IEnumerable<T>> getChildren,
            bool preOrder,
            Func<T, int, TResult> visitor) {

            var postOrder = !preOrder;
            var parentMarker = new object();

            var stack = new Stack<object>();
            stack.Push(node);

            int depth = 0;

            while (stack.Count > 0) {

                object n = stack.Pop();

                if (n == parentMarker) {
                    depth--;
                    n = stack.Pop();
                    if (postOrder)
                        yield return visitor((T)n, depth);
                    continue;
                }

                if (preOrder)
                    yield return visitor((T)n, depth);

                var children = getChildren((T)n);
                if (children != null && children.Any()) {
                    stack.Push(n);
                    stack.Push(parentMarker);
                    stack.PushRange(children.ToEnumerableOfObject().Reverse());
                    depth++;
                } else if (postOrder) {
                    yield return visitor((T)n, depth);
                }
            }
        }

        public static IEnumerable<IGrouping<TParent, TChild>> BreadthFirst<T, TChild, TParent>(
            T node,
            Func<T, IEnumerable<T>> getChildren,
            Func<T, bool> traverseChildren,
            Func<T, bool> shouldReturn,
            Func<T, TChild> childSelector,
            Func<TChild, TParent> parentSelector) {

            TParent parent = default(TParent);
            var parentMarker = new object();
            var siblings = new List<TChild>();

            var queue = new Queue<object>();
            queue.Enqueue(node); // nodes go in as T

            while (queue.Count > 0) {

                var n = queue.Dequeue();

                if (n == parentMarker) {

                    yield return siblings.GroupBy(o => parent).First();
                    parent = (TParent)queue.Dequeue();

                } else {

                    node = (T)n;

                    TChild nodeAsChild = default(TChild);
                    if (shouldReturn(node)) {

                        // nodes are returned as TChildren
                        nodeAsChild = childSelector(node);
                        siblings.Add(nodeAsChild);
                    }

                    if (!traverseChildren(node))
                        continue;

                    var children = getChildren(node);
                    if (children == null || children.None())
                        continue;

                    queue.Enqueue(parentMarker);
                    // nodes which were once children that themselves have children are re-enqueued as parents
                    queue.Enqueue(parentSelector(nodeAsChild));
                    queue.EnqueueRange(getChildren(node).Cast<object>());
                }
            }

            yield return siblings.GroupBy(o => parent).First();
        }

        public static IEnumerable<IGrouping<TParent, TChild>> DepthFirst<T, TChild, TParent>(
            T node,
            Func<T, IEnumerable<T>> getChildren,
            Func<T, bool> traverseChildren,
            Func<T, bool, bool> shouldReturn,
            Func<T, bool, TChild> childSelector,
            Func<TChild, TParent> parentSelector) {

            return DepthFirst(
                node,
                getChildren,
                traverseChildren,
                shouldReturn,
                (o, c, e) => childSelector(o, c),
                parentSelector).ToArray();
        }

        public static IEnumerable<IGrouping<T, T>> DepthFirst<T>(
            T node,
            Func<T, IEnumerable<T>> getChildren) {

            return DepthFirst(
                node,
                getChildren,
                o => true,
               (o, c) => true,
               (o, c) => o,
                o => o);
        }

        private class Grouping<TKey, T> : IGrouping<TKey, T> {

            #region Private Data Members
            private TKey m_key;
            private IEnumerable<T> m_enumerable;
            #endregion

            #region Constructor
            public Grouping(TKey key, IEnumerable<T> enumerable) {
                m_key = key;
                m_enumerable = enumerable ?? new T[] { };
            }
            #endregion

            #region IGrouping<TKey,T> Members
            public TKey Key {
                get { return m_key; }
            }
            #endregion

            #region IEnumerable<T> Members
            public IEnumerator<T> GetEnumerator() {
                return m_enumerable.GetEnumerator();
            }
            IEnumerator IEnumerable.GetEnumerator() {
                return GetEnumerator();
            }
            #endregion
        }

        public static IEnumerable<IGrouping<TParent, TChild>> DepthFirst<T, TChild, TParent>(
            T node,
            Func<T, IEnumerable<T>> getChildren,
            Func<T, bool> traverseChildren,
            Func<T, bool, bool> shouldReturn,
            Func<T, bool, IEnumerable<TChild>, TChild> childSelector,
            Func<TChild, TParent> parentSelector) {

            var parentMarker = new object();

            // the stack contains nodes, parent markers, and a null uber-root
            var stack = new Stack<object>();
            stack.Push(null);                                   // add null as uber-root and...
            stack.Push(parentMarker);                           // ...say uber-root is a parent of...
            stack.Push(node);                                   // ...the root node 

            // the sibling stack accumulates the siblings of each parent node
            var siblingsStack = new Stack<List<TChild>>();
            siblingsStack.Push(null);

            while (stack.Count > 0) {

                object n = stack.Pop();

                // if we finished a set of siblings
                if (n == parentMarker) {

                    // get the siblings
                    var siblings = siblingsStack.Pop();

                    // get the parent
                    var parent = stack.Pop();

                    // if the parent is null then it's the "parent" of the root node
                    if (parent == null) {

                        if (siblings != null && siblings.Any())
                            yield return new Grouping<TParent, TChild>(default(TParent), siblings);

                        // done!
                        break;
                    }

                    // set this node to the parent we just pulled out of the stack
                    node = (T)parent;

                    var nodeHasChildren = siblings != null && siblings.Any();
                    if (shouldReturn(node, nodeHasChildren)) {

                        // this parent will be pushed back onto the stack as a child 
                        var parentAsChild = childSelector(node, true, siblings);

                        // if there were any siblings then yield the parent and it's children
                        yield return new Grouping<TParent, TChild>(parentSelector(parentAsChild), siblings);

                        // if this is the first sibling to be returned cook up a list
                        if (siblingsStack.Peek() == null)
                            // replace the null at the top of the stack with a new list
                            siblingsStack.PopPush(new List<TChild>());

                        // add the node to the list of siblings
                        siblingsStack.Peek().Add(parentAsChild);
                    }

                    continue;
                }

                // if it's not the parent marker than it must be a T node
                node = (T)n;

                if (traverseChildren(node)) {

                    // get any children
                    var children = getChildren(node).Cast<object>().ToArray();

                    // if there are children then delay processing this node until
                    // we know if we want to return any of the children. 
                    if (children != null && children.Any()) {
                        stack.Push(node);
                        stack.Push(parentMarker);
                        stack.PushRange(children.Reverse());
                        siblingsStack.Push(null);
                        continue;
                    }
                }

                if (shouldReturn(node, false)) {

                    // if this is the first sibling to be returned cook up a list
                    if (siblingsStack.Peek() == null)
                        // replace the null at the top of the stack with a new list
                        siblingsStack.PopPush(new List<TChild>());

                    var emptySiblings = new TChild[] { };
                    var leaf = childSelector(node, false, emptySiblings);
                    //yield return new Grouping<TParent, TChild>(parentSelector(leaf), emptySiblings);

                    // add the node to the list of siblings
                    siblingsStack.Peek().Add(leaf);
                }
            }
        }

        public static IEnumerable<IGrouping<TParent, TChild>> LeftTraverse<T, TChild, TParent>(
            T left,
            T right,
            T target,
            Func<T, IEnumerable<T>> getChildren,
            Func<T, bool> hasValue,
            Func<T, T> rightDefaultSelector,
            Func<T, T> targetDefaultSelector,
            Func<T, T, bool> traverseChildren,
            Func<T, T, T, bool, bool> shouldReturn,
            Func<T, T, T, bool, TChild> childSelector,
            Func<TChild, TParent> parentSelector) {

            return DepthFirst(
                new { Left = left, Right = right, Target = target },
                o => getChildren(o.Left)                    // getChildren; get left children 
                    .LeftZip(                               // left zip (take all non-null left nodes)
                        getChildren(o.Right),               // get right children
                        getChildren(o.Target),              // get target children
                        v => hasValue(v),                   // the test to see if the left node has a value
                        rightDefaultSelector,               // if right is null then select a default
                        targetDefaultSelector,              // if target is null then select a default
                        (l, r, t) => new {                  // cook the node 
                            Left = l,
                            Right = r,
                            Target = t
                        }),
                o => traverseChildren(o.Left, o.Right),
               (o, c) => shouldReturn(o.Left, o.Right, o.Target, c),
               (o, c) => childSelector(o.Left, o.Right, o.Target, c),
                o => parentSelector(o)).ToArray();
        }

        public static IEnumerable<IGrouping<TParent, TChild>> LeftTraverse<T, TChild, TParent>(
            T left,
            T right,
            Func<T, IEnumerable<T>> getChildren,
            Func<T, bool> hasValue,
            Func<T, T> defaultSelector,
            Func<T, T, bool> traverseChildren,
            Func<T, T, bool, bool> shouldReturn,
            Func<T, T, bool, TChild> childSelector,
            Func<TChild, TParent> parentSelector) {

            return DepthFirst(
                new { Left = left, Right = right },
                o => getChildren(o.Left)                    // getChildren; get left children 
                    .LeftZip(                               // left zip (take all non-null left nodes)
                        getChildren(o.Right),               // get right children
                        v => hasValue(v),                   // the test to see if the left node has a value
                        defaultSelector,                    // if right is null then select a default
                        (l, r) => new {                     // cook the node 
                            Left = l,
                            Right = r
                        }),
                o => traverseChildren(o.Left, o.Right),
               (o, c) => shouldReturn(o.Left, o.Right, c),  // don't care if siblings are present in a traverse
               (o, c) => childSelector(o.Left, o.Right, c),
                o => parentSelector(o)).ToArray();
        }

        public static IEnumerable<TResult> LeftWhere<T, TResult>(
            T left,
            T right,
            Func<T, IEnumerable<T>> getChildren,
            Func<T, bool> hasValue,
            Func<T, T> defaultSelector,
            Func<T, T, bool> traverseChildren,
            Func<T, T, bool, bool> wherePredicate,
            Func<T, T, TResult> resultSelector) {

            // traverse the left and right children in parallel
            foreach (var pair in TreeTraverse.LeftTraverse(
                left, right,
                getChildren,
                hasValue,
                defaultSelector,
                traverseChildren,
                wherePredicate,
                (l, r, c) => new { Left = l, Right = r },   // childSelector
                o => default(object))                       // parentSelector; don't care 
                .Flatten().ToArray()) {                     // flatten the siblings into decendents

                // select out the left and right nodes
                var leftNode = pair.Left;
                var rightNode = pair.Right;

                // generate the result
                yield return resultSelector(leftNode, rightNode);
            }
        }
    }
}
